// Copyright 2022 jeff.li. and/or its affiliates.

#pragma once

#include <memory>
#include <string>
#include <thread>
#include <unordered_map>

#include <tbir/pipeline/tx_session.h>
#include <tbir/runtime/runtime_value.h>
#include <tbir/runtime/threadpool/mpmc_bounded_queue.h>

namespace tbir::runtime::server {

    class SimpleMPMCServer {
        class Runnable {
        public:
            Runnable(const std::unordered_map<std::string, RTValue> *inputs,
                     std::vector<std::pair<std::string, RTValue>> *outputs)
                    : inputs(inputs), outputs(outputs) {};

            virtual void Run(const TXSession *sess_ptr) {
                try {
                    *outputs = sess_ptr->Run(*inputs);
                } catch (...) {
                    throw_exception_ = true;
                    except_ptr_ = std::current_exception();
                }
            }

            void SetDone() {
                finish_ = true;
            }

            bool Done() const {
                return finish_;
            }

            bool HasException() const {
                return throw_exception_;
            }

            std::exception_ptr ExceptionPtr() const {
                return except_ptr_;
            }

        private:
            volatile bool finish_ = false;
            const std::unordered_map<std::string, RTValue> *inputs;
            std::vector<std::pair<std::string, RTValue>> *outputs;

            bool throw_exception_ = false;
            std::exception_ptr except_ptr_ = nullptr;
        };

        using RunnablePtr = std::shared_ptr<Runnable>;

        struct RunnableWithTimeCost : public Runnable {
            RunnableWithTimeCost(const std::unordered_map<std::string, RTValue> *inputs,
                                 std::vector<std::pair<std::string, RTValue>> *outputs)
                    : Runnable(inputs, outputs) {
            }

            uint64_t time_cost = 0;

            void Run(const TXSession *sess_ptr) override {
                auto begin = EnvTime::Default()->NowMicros();
                Runnable::Run(sess_ptr);
                auto end = EnvTime::Default()->NowMicros();
                time_cost = end - begin;
            }
        };

    public:
        explicit SimpleMPMCServer(std::vector<std::shared_ptr<TXSession>> handlers,
                                  std::string name_prefix);

        ~SimpleMPMCServer() {
            stop();
        }

        void start();

        void stop();

        std::vector<std::pair<std::string, RTValue>> process(
                const std::unordered_map<std::string, RTValue> &feed_dict);

        std::vector<std::pair<std::string, RTValue>> process_with_tc(
                const std::unordered_map<std::string, RTValue> &feed_dict,
                uint64_t *enqueue_tc_us,
                uint64_t *real_run_tc_us);

    protected:
        static void ThreadEntry(SimpleMPMCServer *pool, TXSession *sess_ptr, const std::string &name);

    private:
        // need to keep track of threads so we can join them
        std::vector<std::thread> workers_;
        std::vector<std::shared_ptr<TXSession>> handlers_;
        // the task queue
        ::tbir::runtime::MPMCBoundedQueue<RunnablePtr> tasks_;
        // stop flag
        bool stop_ = false;
        std::string name_;
    };

}  // namespace tbir::runtime::server
